package com.zzh.h265toupin

import android.content.res.Resources
import android.hardware.display.DisplayManager
import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.media.projection.MediaProjection
import android.util.Log
import com.zzh.h265toupin.sockets.SocketPushServer
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import java.nio.ByteBuffer

class H265Encode(
    private val socketServer: SocketPushServer,
    private val mediaProjection: MediaProjection
) {
    private val dm = Resources.getSystem().displayMetrics
    private val screenWidth = dm.widthPixels
    private val screenHeight = dm.heightPixels

    private lateinit var mediaCodec: MediaCodec

    private var startFlag = true
    private var totalSize = 0


    fun startPush() {
        val width = if (screenWidth % 2 == 0) screenWidth else screenWidth - 1
        val height = if (screenHeight % 2 == 0) screenHeight else screenHeight - 1
        Log.d("TAG_H265Encode", "屏幕分辨率 screenWidth:$screenWidth screenHeight:$screenHeight")
        Log.d("TAG_H265Encode", "编码分辨率 width:$width height:$height")
        //创建h264编码器
        mediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_HEVC)
       // mediaCodec.reset()
        val format =
            MediaFormat.createVideoFormat(
                MediaFormat.MIMETYPE_VIDEO_HEVC,
                width,
                height
            ).apply {
                //配置编码帧率
                setInteger(MediaFormat.KEY_FRAME_RATE, 15)
                //配置编码码率1000kb/s
                setInteger(MediaFormat.KEY_BIT_RATE, 1000_000)
                //配置关键帧I帧间隔为2秒一个
                setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 2)
                //配置编码颜色格式，YUV420/NV21等，这里录制的屏幕数据是graphicbuffer元数据使用COLOR_FormatSurface
                setInteger(
                    MediaFormat.KEY_COLOR_FORMAT,
                    MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface
                )
            }
        mediaCodec.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        //通过MediaCodec提供surface,让屏幕录制的数据输入到MediaCodec
        val meidaSurface = mediaCodec.createInputSurface()
        CoroutineScope(Dispatchers.IO).launch {
            mediaCodec.start()
            //让mediaProjection输出到MediaCodec
            mediaProjection.createVirtualDisplay(
                "screen-codec",
                screenWidth, //输出录屏宽大小
                screenHeight, //输出录屏高大小
                1,  //输出录屏dpi数，不放大缩小就直接1
                DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC, //输出到虚拟的画布标记
                meidaSurface, //MeidaCodec画布
                null,
                null
            )

            //创建一个存放编码的数据缓冲区
            try {
                val bufferInfo = MediaCodec.BufferInfo()
                while (startFlag) {
                    //100毫秒(100_000微妙)内完成编码的数据，放到bufferInfo内，并返回编码完成的通道编号,没有完成会返回-1
                    val channal = mediaCodec.dequeueOutputBuffer(bufferInfo, 100_000)
                    if (channal >= 0) {
                        //获取对应通道编码完成的字节数据
                        val buffer: ByteBuffer? = mediaCodec.getOutputBuffer(channal)
                        buffer?.let { dealPushFrame(it, bufferInfo) }
                        totalSize += bufferInfo.size
                        Log.i("TAG_MainActivity", "已编码：${totalSize / 1024 / 1024f}MB")
                        //编码完成后释放编码完成的通道
                        mediaCodec.releaseOutputBuffer(channal, false)
                    }
                }
            } catch (e: Exception) {
                e.printStackTrace()
            }
        }
    }

    fun stop() {
        startFlag = false
        mediaProjection.stop()
        mediaCodec.release()
    }

    //H265的第一帧I帧类型
    private val NAL_I = 19

    //H265的VPS类型
    private val NAL_VPS = 32

    //记录第一帧信息帧
    private lateinit var vps_sps_pps_buf: ByteArray

    //记录最新的I帧，用于一连接则推送I帧
    lateinit var curIFrame: ByteArray

    //处理推流的帧
    private fun dealPushFrame(buffer: ByteBuffer, bufferInfo: MediaCodec.BufferInfo) {
        //查找分隔符是00 00 01还是 00 00 00 01，并跳到分隔符后面开始的字节取出第一帧信息
        val offset = if (buffer.get(2).toInt() == 0x01) 3 else 4
        //分析H265第一帧的帧类型, 帧类型通过 与上 01111110（0x7E）并右移1位获取
        val type = (buffer.get(offset).toInt() and 0x7E) shr 1
        val screenFrame = ByteArray(bufferInfo.size).apply { buffer.get(this) }
        when (type) {
            NAL_VPS -> {
                //记住开始编码的第一帧VPS/SPS/PPS数据，用于每次遇到I帧就拼接上一齐发送，VPS/SPS/PPS数据只会出现一次
                vps_sps_pps_buf = screenFrame
                Log.v("TAG_Encode", "VPS-SPS-PPS帧数据  ${vps_sps_pps_buf.contentToString()}")
            }
            NAL_I -> {
                //遇到I帧编码时，每个I帧都带上VPS/SPS/PPS数据再发送
                val newIFrame = ByteArray(vps_sps_pps_buf.size + screenFrame.size)
                System.arraycopy(vps_sps_pps_buf, 0, newIFrame, 0, vps_sps_pps_buf.size)
                System.arraycopy(screenFrame, 0, newIFrame, vps_sps_pps_buf.size, screenFrame.size)
                curIFrame = newIFrame
                socketServer.sendData(newIFrame)
                Log.v("TAG_Encode", "I帧数据  ${newIFrame.contentToString()}")
            }
            else -> {
                //P帧B帧就直接发送
                socketServer.sendData(screenFrame)
                Log.v("TAG_Encode", "PB帧数据  ${screenFrame.contentToString()}")
            }
        }

    }

}